open Base

type t =
  | Arm of Draw_tree.dir * t list
  | Figure of fig_b
  | With_handler of Event.handler * t
  | With_fold of Source.value
                 * (Event.t -> Source.value -> Source.value)
                 * (Source.t -> t)
and fig_b =
  Draw_tree.fig Behavior.t

let arm dir ts =
  Arm(dir, ts)
let rect b =
  let f (dim, c) = Draw_tree.Rect(dim, c) in
  Figure(Behavior.map b ~f)
let callback ~f t =
  let handler e = f e ; Source.empty_map in
  With_handler(handler, t)
let fold ~init ~f make_child =
  With_fold (init, f, fun src -> make_child (Behavior.of_source src))

type sink = Draw_tree.Path.t * fig_b

type callbacks =
  { sinks : sink list Source.map
  ; handlers : Event.handler Event.map
  }

let empty_callbacks =
  { sinks = Source.empty_map
  ; handlers = Event.empty_map }

let mount state t0 : Source.State.t * callbacks * Draw_tree.t =
  let state = ref state in
  let rec mount rev_path callbacks = function
    | Arm(dir, ts) ->
       let mount' i cbs t = mount (i::rev_path) cbs t in
       let callbacks, dts = List.fold_mapi ts ~init:callbacks ~f:mount' in
       callbacks, Draw_tree.arm dir dts

    | Figure(fb) ->
       let deps = Behavior.dependencies fb in
       let callbacks =
         if Sequence.is_empty deps then
           callbacks
         else
           let path = Draw_tree.Path.of_list_rev rev_path in
           let register cbs src =
             { cbs with
               sinks = Map.add_multi cbs.sinks
                         ~key:src ~data:(path, fb) }
           in
           Sequence.fold deps ~init:callbacks ~f:register
       in
       callbacks, Draw_tree.figure (Behavior.sample fb !state)

    | With_handler(handler, child) ->
       mount_handler rev_path callbacks handler child

    | With_fold(init, f, make_child) ->
       let src = Source.create () in
       let () = state := Source.State.set src init !state in
       mount_handler rev_path callbacks
         (fun ev -> Map.singleton (module Source) src (f ev))
         (make_child src)

  and mount_handler rev_path callbacks handler child =
    let ev_id = Event.Id.gen () in
    let callbacks, dt = mount (0::rev_path) callbacks child in
    { callbacks with
      handlers = Map.add_exn callbacks.handlers
                   ~key:ev_id ~data:handler },
    Draw_tree.capture ev_id dt

  in
  let cbs, dt = mount [] empty_callbacks t0 in
  !state, cbs, dt
